package org.fujay.langchain4jserver.infrastructure.persistence.repository;

import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import lombok.RequiredArgsConstructor;
import org.fujay.langchain4jserver.domain.conversation.entity.MessageEntity;
import org.fujay.langchain4jserver.domain.conversation.repository.MessageRepository;
import org.fujay.langchain4jserver.domain.conversation.vo.ConversationId;
import org.fujay.langchain4jserver.domain.conversation.vo.MessageId;
import org.fujay.langchain4jserver.infrastructure.persistence.assembler.MessageAssembler;
import org.fujay.langchain4jserver.infrastructure.persistence.mybatis.entity.MessagePo;
import org.fujay.langchain4jserver.infrastructure.persistence.mybatis.mapper.MessageMapper;
import org.springframework.stereotype.Repository;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.util.CollectionUtils;

import java.util.List;
import java.util.Optional;

@Repository
@RequiredArgsConstructor
public class MessageRepositoryImpl implements MessageRepository {

    private final MessageMapper messageMapper;
    private final MessageAssembler messageAssembler;

    @Override
    public Optional<MessageEntity> findById(MessageId id) {
        MessagePo po = messageMapper.selectById(id.value());
        return Optional.ofNullable(messageAssembler.toDomainEntity(po));
    }

    @Override
    public List<MessageEntity> findByConversationId(ConversationId conversationId) {
        LambdaQueryWrapper<MessagePo> queryWrapper =
                Wrappers.<MessagePo>lambdaQuery()
                        .eq(MessagePo::getConversationId, conversationId.value())
                        .orderByAsc(MessagePo::getTimestamp); // Order by time
        List<MessagePo> pos = messageMapper.selectList(queryWrapper);
        return messageAssembler.toDomainEntities(pos);
    }

    @Override
    public void save(MessageEntity message, ConversationId conversationId) {
        MessagePo po = messageAssembler.toPO(message);
        po.setConversationId(conversationId.value());
        // Use saveOrUpdate logic based on ID presence or isNew flag
        if (po.getId() != null && messageMapper.selectById(po.getId()) != null) {
            messageMapper.updateById(po);
        } else {
            // Ensure ID is set if new (assuming Assembler doesn't set if null)
            if (po.getId() == null) {
                po.setId(message.getId().value());
            }
            messageMapper.insert(po);
        }
    }

    @Override
    @Transactional(rollbackFor = Exception.class)
    public void saveAll(List<MessageEntity> messages, ConversationId conversationId) {
        if (CollectionUtils.isEmpty(messages)) {
            return;
        }
        // Consider using batch insert/update for performance
        messages.forEach(message -> save(message, conversationId));
    }

    @Override
    @Transactional(rollbackFor = Exception.class)
    public void deleteById(MessageId id) {
        messageMapper.deleteById(id.value());
    }

    @Override
    @Transactional(rollbackFor = Exception.class)
    public void deleteByConversationId(ConversationId conversationId) {
        LambdaQueryWrapper<MessagePo> deleteWrapper =
                Wrappers.<MessagePo>lambdaQuery().eq(MessagePo::getConversationId, conversationId.value());
        messageMapper.delete(deleteWrapper);
    }

    @Override
    public MessageId nextId() {
        // MessageId expects String, generate UUID
        return MessageId.create();
    }
}
